feat: Refactor LLM model zoo and add KV cache support #3527

peri044 · 2025-05-20T23:30:28Z

Description

This PR redesigns our LLM model compilation, unifies it, fixes output mismatch and performance issues. This PR also implements KV caching using native TensorRT.

Fixes # (issue)

Type of change

New feature (non-breaking change which adds functionality)
This change requires a documentation update

Checklist:

My code follows the style guidelines of this project (You can use the linters)
I have performed a self-review of my own code
I have commented my code, particularly in hard-to-understand areas and hacks
I have made corresponding changes to the documentation
I have added tests to verify my fix or my feature
New and existing unit tests pass locally with my changes
I have added the relevant labels to my PR in so that relevant reviewers are notified

…s landed

Signed-off-by: Dheeraj Peri <[email protected]>

tools/llm/cache_utils.py

tools/llm/sdpa_converter.py

narendasan · 2025-06-13T21:04:44Z

tools/llm/test_qwen2.5_components.py

@@ -0,0 +1,193 @@
+import torch


Should we create a tests/py/tools/llm?

We can. These tests helped me debug accuracy/functionality issues during the development. It would take some time to get them working in CI. I can add them as a part of cherry-pick PR after this. What do you think ?

tools/llm/cache_utils.py

zewenli98

Mostly LGTM. just some questions

zewenli98 · 2025-06-14T00:45:57Z

py/torch_tensorrt/dynamo/_compiler.py

@@ -837,6 +859,7 @@ def contains_metadata(gm: torch.fx.GraphModule) -> bool:
            continue
        submodule_node_dict[node.name] = node

+    preserve_module_specs(original_in_spec, original_out_spec, partitioned_module)


Should set spec for partitioned_module's children?

The spec for children is actually set by the partitioning. I'm only trying to preserve the output spec for the entire partitioned_module (It should be the same as original module unless we change it explicitly which is the case in KV caching)

zewenli98 · 2025-06-14T05:24:03Z

examples/dynamo/weight_streaming_example.py

-from utils import export_llm
+
+
+def export_llm(model, inputs, min_seq_len=1, max_seq_len=16):


why not directly call tools/llm/utils.py?

We have to use sys.append to include the path to tools/llm/utils.py to do that. I wanted this example to be standalone so I just copied this function.

zewenli98

same questions above

narendasan

LGTM only suggestion is moving tests into a tests subdirectory or into the tests directory. Is optional.

narendasan · 2025-06-17T02:56:00Z

tools/llm/utils.py

+    return timings
+
+
+def recordStats(backend, timings, precision, batch_size=1, compile_time_s=None):


Super minor but everything else is snake case

peri044 and others added 9 commits April 25, 2025 21:03

chore: kv cache prototyping

e5723db

chore: add sdpa converter/lowering

3e0b46a

feat: implement static/dynamic kv cache in Torch-TRT

e30fa42

chore: updates

a3a202f

chore: updates

3642614

chore: updates

3688630

chore: refactor updates

c9f5f27

chore: refactor updates

0cb0dcc

chore: updates

6cbb1bd

peri044 added the WIP Work is in progress, pull request should not be merged yet label May 20, 2025

facebook-github-bot added the cla signed label May 20, 2025

peri044 marked this pull request as draft May 20, 2025 23:30

github-actions bot requested a review from gs-olive May 20, 2025 23:30

peri044 changed the title ~~feat : caching attempts~~ feat: caching attempts May 20, 2025

peri044 removed the request for review from gs-olive May 21, 2025 00:40

Chengzhe Xu and others added 3 commits May 27, 2025 19:13

chore: updates

f539b55

chore: updates

0dc3a7e

chore: rebase with main

095b5cf

This comment was marked as spam.

Sign in to view

This comment was marked as outdated.

Sign in to view

peri044 added 2 commits June 4, 2025 09:15

feat: Refactor LLM runner and implemented support for Qwen family

600e363

chore: updates

9309725

peri044 force-pushed the kv_cache branch from 19bc80b to 9309725 Compare June 4, 2025 09:17

This comment was marked as outdated.

Sign in to view

chore: updates

a50e7ac

chore: remove is_causal input now that causal attention enhancement i…

c1f0053

…s landed

This comment was marked as spam.

Sign in to view

chore: refactor

8301ee6

This comment was marked as spam.

Sign in to view

peri044 removed WIP Work is in progress, pull request should not be merged yet labels Jun 12, 2025

peri044 marked this pull request as ready for review June 12, 2025 22:36

This comment was marked as spam.

Sign in to view

chore: move code to tools/llm

ecf88d1

This comment was marked as spam.

Sign in to view

peri044 added 4 commits June 12, 2025 22:55

chore: remove files not considered for release

723d1b2

Signed-off-by: Dheeraj Peri <[email protected]>

chore: rebase with main

b835ae9

chore: updates

5eabf65

chore: Add README.md

806616d

peri044 changed the title ~~feat: caching attempts~~ feat: Refactor LLM optimization and add KV cache support Jun 13, 2025

peri044 changed the title ~~feat: Refactor LLM optimization and add KV cache support~~ feat: Refactor LLM model zoo and add KV cache support Jun 13, 2025

chore: rebase with main

53a2c56

github-actions bot removed the component: lowering Issues re: The lowering / preprocessing passes label Jun 13, 2025

chore: add docs

c2cbd5a

github-actions bot added the documentation Improvements or additions to documentation label Jun 13, 2025

peri044 added 2 commits June 13, 2025 13:31

chore: Add a tutorial

81114f8

chore: fix model name

57da513

narendasan reviewed Jun 13, 2025

View reviewed changes

tools/llm/cache_utils.py Outdated Show resolved Hide resolved

narendasan reviewed Jun 13, 2025

View reviewed changes

tools/llm/sdpa_converter.py Outdated Show resolved Hide resolved

narendasan reviewed Jun 13, 2025

View reviewed changes

tools/llm/cache_utils.py Show resolved Hide resolved

zewenli98 reviewed Jun 14, 2025

View reviewed changes

chore: address review comments

deec75f

zewenli98 reviewed Jun 17, 2025

View reviewed changes

narendasan approved these changes Jun 17, 2025

View reviewed changes

		from utils import export_llm


		def export_llm(model, inputs, min_seq_len=1, max_seq_len=16):

		return timings


		def recordStats(backend, timings, precision, batch_size=1, compile_time_s=None):

feat: Refactor LLM model zoo and add KV cache support #3527

Are you sure you want to change the base?

feat: Refactor LLM model zoo and add KV cache support #3527

Uh oh!

Conversation

peri044 commented May 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Checklist:

Uh oh!

This comment was marked as spam.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as spam.

Uh oh!

This comment was marked as spam.

Uh oh!

This comment was marked as spam.

Uh oh!

This comment was marked as spam.

Uh oh!

Uh oh!

Uh oh!

narendasan Jun 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

peri044 Jun 13, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

zewenli98 left a comment

Choose a reason for hiding this comment

Uh oh!

zewenli98 Jun 14, 2025

Choose a reason for hiding this comment

Uh oh!

peri044 Jun 17, 2025

Choose a reason for hiding this comment

Uh oh!

zewenli98 Jun 14, 2025

Choose a reason for hiding this comment

Uh oh!

peri044 Jun 16, 2025

Choose a reason for hiding this comment

Uh oh!

zewenli98 left a comment

Choose a reason for hiding this comment

Uh oh!

narendasan left a comment

Choose a reason for hiding this comment

Uh oh!

narendasan Jun 17, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

peri044 commented May 20, 2025 •

edited

Loading

narendasan Jun 13, 2025 •

edited

Loading